# -*- coding: utf-8 -*-
"""
@author: administrator
"""
import torch.hub
import torchvision.models
from PIL import Image

from utils import pytorch_utils as utils
from utils.assets import Assets

if __name__ == '__main__':
    net = torch.hub.load(Assets.pytorch_repo, "resnet50", weights=torchvision.models.ResNet50_Weights)
    net.eval()

    print("dog:", Assets.img_dog)

    transform = torchvision.transforms.Compose([
        torchvision.transforms.Resize(225),
        torchvision.transforms.CenterCrop(224),
        torchvision.transforms.ToTensor(),
        torchvision.transforms.Normalize(mean=[0.485, 0.456, 0.406], std=[0.229, 0.224, 0.225])
    ])
    image = Image.open(Assets.img_dog)
    image = transform(image)
    image = image.unsqueeze(0)

    device = utils.get_device()
    image = image.to(device)

    with torch.no_grad():
        net.to(device)
        out = net(image)

    pro = torch.nn.functional.softmax(out[0], dim=0)

    with open("../../assets/imagenet_classes.txt", "r") as f:
        categories = [s.strip() for s in f.readlines()]
    # Show top categories per image
    top5_prob, top5_catid = torch.topk(pro, 5)
    for i in range(top5_prob.size(0)):
        print(categories[top5_catid[i]], top5_prob[i].item())
